import ast
import logging
from typing import Dict, List, Set

from black import format_str, FileMode

from 程序核心代码.ast节点.节点定义.基础定义.节点异常定义 import 节点异常

logger = logging.getLogger(__name__)

class 代码格式化器:
    """代码格式化器类，负责对生成的源代码进行格式化"""
    
    def __init__(self):
        self.行长度限制 = 88  # black默认行长度
        self.缩进字符 = "    "
        self.换行符 = "\n"
        self.导入排序规则 = {
            "标准库": 0,
            "第三方库": 1,
            "本地模块": 2
        }
        
    def 格式化代码(self, 源代码: str) -> str:
        """格式化Python源代码"""
        try:
            # 使用black进行基本格式化
            格式化后代码 = format_str(源代码, mode=FileMode(line_length=self.行长度限制))
            
            # 优化导入语句
            格式化后代码 = self._优化导入(格式化后代码)
            
            return 格式化后代码
            
        except Exception as e:
            logger.error(f"代码格式化失败: {str(e)}")
            raise 节点异常(f"代码格式化失败: {str(e)}") from e
            
    def _优化导入(self, 源代码: str) -> str:
        """优化导入语句的顺序和格式"""
        try:
            树 = ast.parse(源代码)
            导入语句 = []
            其他语句 = []
            
            for 节点 in ast.iter_child_nodes(树):
                if isinstance(节点, (ast.Import, ast.ImportFrom)):
                    导入语句.append(节点)
                else:
                    其他语句.append(节点)
                    
            # 对导入语句进行分类和排序
            分类导入 = self._分类导入语句(导入语句)
            排序导入 = self._排序导入语句(分类导入)
            
            # 重建源代码
            新代码 = []
            for 导入组 in 排序导入:
                for 导入 in 导入组:
                    新代码.append(ast.unparse(导入))
                新代码.append("")  # 添加空行分隔不同组的导入
                
            # 添加其他语句
            for 语句 in 其他语句:
                新代码.append(ast.unparse(语句))
                
            return self.换行符.join(新代码)
            
        except Exception as e:
            logger.warning(f"导入优化失败: {str(e)}")
            return 源代码
            
    def _分类导入语句(self, 导入语句: List[ast.AST]) -> Dict[str, List[ast.AST]]:
        """将导入语句分类为标准库、第三方库和本地模块"""
        分类 = {
            "标准库": [],
            "第三方库": [],
            "本地模块": []
        }
        
        标准库集合 = set(self._获取标准库列表())
        
        for 导入 in 导入语句:
            if isinstance(导入, ast.Import):
                模块名 = 导入.names[0].name.split(".")[0]
            else:  # ImportFrom
                模块名 = 导入.module.split(".")[0] if 导入.module else ""
                
            if 模块名 in 标准库集合:
                分类["标准库"].append(导入)
            elif "." in 模块名 or 模块名.startswith("_"):
                分类["本地模块"].append(导入)
            else:
                分类["第三方库"].append(导入)
                
        return 分类
        
    def _排序导入语句(self, 分类导入: Dict[str, List[ast.AST]]) -> List[List[ast.AST]]:
        """按规则对导入语句进行排序"""
        结果 = []
        for 类别 in sorted(分类导入.keys(), key=lambda k: self.导入排序规则[k]):
            if 分类导入[类别]:
                排序后导入 = sorted(分类导入[类别], 
                                key=lambda x: ast.unparse(x).lower())
                结果.append(排序后导入)
        return 结果
        
    def _获取标准库列表(self) -> Set[str]:
        """获取Python标准库列表"""
        import sys
        import distutils.sysconfig as sysconfig
        
        标准库路径 = sysconfig.get_python_lib(standard_lib=True)
        标准库列表 = set()
        
        for 模块名 in sys.stdlib_module_names:
            标准库列表.add(模块名)
            
        return 标准库列表
        
    def 设置行长度限制(self, 长度: int) -> None:
        """设置代码行长度限制"""
        self.行长度限制 = 长度
        
    def 设置缩进(self, 缩进字符: str = "    ") -> None:
        """设置缩进字符"""
        self.缩进字符 = 缩进字符
        
    def 设置换行符(self, 换行符: str = "\n") -> None:
        """设置换行符"""
        self.换行符 = 换行符 